import numpy as np
import pandas as pd
from datetime import datetime, timedelta
import warnings
import os
import matplotlib.pyplot as plt

warnings.filterwarnings('ignore')
import numpy as np
from scipy.stats import truncnorm

# 定义均值和标准差
mu = 0.0  # 均值
sigma = 0.05  # 标准差（调整这个值可以控制分布的宽度）

# 定义截断范围 [a, b]
a = -0.1
b = 0.1

# 计算截断正态分布的参数
a_norm = (a - mu) / sigma
b_norm = (b - mu) / sigma



# 1. 数据生成和保存
def get_raw_data(stock_num=50, days=1000, save_path='factor_data.xlsx'):
    """
    生成随机股票数据
    :param stock_num: 股票数量
    :param days: 交易日数量
    :param save_path: 保存路径
    :return: 包含所有数据的DataFrame
    """
    if os.path.exists(save_path):
        return  pd.read_excel(save_path)
    np.random.seed(42)

    # 生成日期
    dates = pd.date_range(end=datetime.today(), periods=days)

    # 生成股票代码
    stocks = [f'STK_{i:03d}' for i in range(1, stock_num + 1)]

    # 生成基础数据
    data = []
    stock_dict = {}
    for stock in stocks:
        stock_dict[stock] = np.random.uniform(10, 100)

    for date in dates:
        for stock in stocks:
            # 生成随机价格和交易量
            open_price = stock_dict[stock]
            close = open_price*(1+truncnorm.rvs(a_norm, b_norm, loc=mu, scale=sigma))
            stock_dict[stock] = close

            high = max(open_price, close) * np.random.uniform(1, 1.1)
            low = min(open_price, close) * np.random.uniform(0.9, 1)
            volume = np.random.randint(10000, 1000000)

            # 生成一些基本面数据
            pe = np.random.uniform(5, 30)
            pb = np.random.uniform(0.5, 5)
            market_cap = np.random.uniform(1e8, 1e10)

            data.append([date, stock, open_price, high, low, close, volume, pe, pb, market_cap])

    # 创建DataFrame
    columns = ['date', 'stock', 'open', 'high', 'low', 'close', 'volume', 'pe', 'pb', 'market_cap']
    df = pd.DataFrame(data, columns=columns)

    # 计算收益率 (未来1天)
    df['return'] = df.groupby('stock')['close'].pct_change().shift(-1)

    # 保存到Excel
    df.to_excel(save_path, index=False)
    print(f"数据已保存到 {save_path}")

    plt.figure(figsize=(12, 6))

    # 按股票分组绘制收盘价
    for stock in df['stock'].unique():
        stock_data = df[df['stock'] == stock]
        plt.plot(stock_data['date'],
                 stock_data['close'],
                 label=stock,  # 可选：显示数据点
                 linestyle='-')

    # 添加图例和标签
    plt.title('Stock Closing Prices Comparison')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.grid(True, linestyle='--', alpha=0.6)  # 网格线
    plt.xticks(rotation=45)  # 旋转日期标签
    plt.tight_layout()  # 避免标签重叠
    plt.show()
    return df